package tom.spring.cloud.gateway;

import com.google.common.collect.Maps;
import lombok.Setter;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.BeansException;
import org.springframework.beans.factory.DisposableBean;
import org.springframework.boot.SpringApplication;
import org.springframework.boot.autoconfigure.SpringBootApplication;
import org.springframework.cloud.client.discovery.EnableDiscoveryClient;
import org.springframework.cloud.gateway.filter.ratelimit.AbstractRateLimiter;
import org.springframework.cloud.gateway.filter.ratelimit.KeyResolver;
import org.springframework.cloud.gateway.filter.ratelimit.RateLimiter;
import org.springframework.cloud.gateway.route.RouteDefinitionRouteLocator;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationContextAware;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.ComponentScan;
import org.springframework.context.annotation.FilterType;
import org.springframework.context.annotation.Primary;
import org.springframework.security.config.annotation.web.reactive.EnableWebFluxSecurity;
import org.springframework.security.config.web.server.ServerHttpSecurity;
import org.springframework.security.web.server.SecurityWebFilterChain;
import org.springframework.util.Assert;
import org.springframework.validation.Validator;
import org.springframework.validation.annotation.Validated;
import org.springframework.web.reactive.config.EnableWebFlux;
import reactor.core.publisher.Mono;

import javax.validation.constraints.Min;
import javax.validation.constraints.NotNull;
import java.lang.ref.SoftReference;
import java.net.InetSocketAddress;
import java.util.Calendar;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ScheduledThreadPoolExecutor;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicLong;

/**
 * 基于配置的gateway, profiles: ratelimit2
 * http://localhost:8090/actuator/gateway/routes
 * 自定义KeyResolver 自定义rateLimiter
 * {@link org.springframework.cloud.gateway.config.GatewayAutoConfiguration#requestRateLimiterGatewayFilterFactory(RateLimiter, KeyResolver)}
 * @author ZHUFEIFEI
 */
@Slf4j
@EnableWebFlux
@EnableDiscoveryClient
@SpringBootApplication
@ComponentScan(excludeFilters = {
        @ComponentScan.Filter(type = FilterType.ASSIGNABLE_TYPE,
                classes = {Gateway01Application.class, Gateway02Application.class, Gateway03Application.class})
})
public class Gateway04Application {

    public static void main(String[] args) {
        SpringApplication.run(Gateway04Application.class, args);
    }

    /**
     * 自定义keyResolver
     * @return
     */
    @Bean
    public KeyResolver myKeyResolver() {
        //自定义根据hostname客户端请求地址做key标识
        return exchange -> {
            //根据客户端地址做限制
            InetSocketAddress addr = exchange.getRequest().getRemoteAddress();
            if (addr != null) {
                log.info("use client address limit : {}", addr.getAddress());
                return Mono.just(addr.getHostString());
            }
            log.info("use request uri limit : {}", exchange.getRequest().getURI());
            //获取不到则根据接口名做限制
            return Mono.just(exchange.getRequest().getURI().toString());
        };
    }

    /**
     * 由于引入了data-redis-reactive依赖,会注入一个redisRateLimiter,所以会出现两个RateLimiter的bean,
     * 使用Primary指定一个主的
     * {@link org.springframework.cloud.gateway.config.GatewayAutoConfiguration#requestRateLimiterGatewayFilterFactory(RateLimiter, KeyResolver)}
     * {@link org.springframework.cloud.gateway.filter.factory.RequestRateLimiterGatewayFilterFactory}
     * @return
     */
    @Primary
    @Bean
    public RateLimiter myRateLimiter() {
        return new MyRateLimiter(5);
    }



    @EnableWebFluxSecurity
    class MyFluxSecurityConfig {

        @Bean
        public SecurityWebFilterChain securityWebFilterChain(ServerHttpSecurity http) {
            //允许所有请求
            return http.authorizeExchange().anyExchange().permitAll().and().build();
        }
    }
}

/**
 * 参考实现 {@link org.springframework.cloud.gateway.filter.ratelimit.RedisRateLimiter} 采用令牌桶算法<br>
 * 自定义实现漏桶算法
 */
@Slf4j
class MyRateLimiter extends AbstractRateLimiter<MyRateLimiter.Config> implements ApplicationContextAware, DisposableBean {

    public static final String CONFIGURATION_PROPERTY_NAME = "my-rate-limiter";
    public static final String REMAINING_HEADER = "X-RateLimit-Remaining";
    public static final String BURST_CAPACITY_HEADER = "X-RateLimit-Burst-Capacity";
    @Setter
    private String remainingHeader = REMAINING_HEADER;
    @Setter
    private String burstCapacityHeader = BURST_CAPACITY_HEADER;

    private Config defaultConfig;
    private AtomicBoolean initialized = new AtomicBoolean(false);
    private QpsLimit limit;

    public MyRateLimiter(int defaultBurstCapacity) {
        super(MyRateLimiter.Config.class, CONFIGURATION_PROPERTY_NAME, null);
        this.defaultConfig = new Config()
                .setBurstCapacity(defaultBurstCapacity);
        this.limit = new QpsLimit();
        this.limit.start();
    }

    @Override
    public Mono<Response> isAllowed(String routeId, String id) {
        if (!this.initialized.get()) {
            throw new IllegalStateException("MyRateLimiter is not initialized");
        }
        if (log.isDebugEnabled()) {
            log.debug("isAllowed : {} ==> {}", routeId, id);
        }
        Config routeConfig = loadConfiguration(routeId);
        int burstCapacity = routeConfig.getBurstCapacity();

        long n = this.limit.limit(id, burstCapacity);
        Response response;
        response = new Response(n > 0, getHeaders(routeConfig, n));
        if (log.isDebugEnabled()) {
            log.debug("response: " + response);
        }
        return Mono.just(response);
    }

    @Override
    public void setApplicationContext(ApplicationContext applicationContext) throws BeansException {
        if (initialized.compareAndSet(false, true)) {
            if (applicationContext.getBeanNamesForType(Validator.class).length > 0) {
                super.setValidator(applicationContext.getBean(Validator.class));
            }
        }
    }

    @NotNull
    public HashMap<String, String> getHeaders(Config config, Long tokensLeft) {
        HashMap<String, String> headers = new HashMap<>();
        headers.put(this.remainingHeader, tokensLeft.toString());
        headers.put(this.burstCapacityHeader, String.valueOf(config.getBurstCapacity()));
        return headers;
    }

    Config loadConfiguration(String routeId) {
        Config routeConfig = getConfig().getOrDefault(routeId, defaultConfig);

        if (routeConfig == null) {
            routeConfig = getConfig().get(RouteDefinitionRouteLocator.DEFAULT_FILTERS);
        }

        if (routeConfig == null) {
            throw new IllegalArgumentException("No Configuration found for route " + routeId +" or defaultFilters");
        }
        return routeConfig;
    }

    @Override
    public void destroy() throws Exception {
        this.limit.stop();
    }

    @Validated
    public static class Config {

        @Min(1)
        private int burstCapacity = 1;

        public int getBurstCapacity() {
            return burstCapacity;
        }

        public MyRateLimiter.Config setBurstCapacity(int burstCapacity) {
            this.burstCapacity = burstCapacity;
            return this;
        }

        @Override
        public String toString() {
            return "Config{" +
                    "burstCapacity=" + burstCapacity +
                    '}';
        }
    }
}

@Slf4j
class QpsLimit {
    /**
     *  内存不足时回收部分对象
     */
    private final Map<String, SoftReference<Map<Integer, AtomicLong>>> limit;
    private ScheduledExecutorService timer;
    private AtomicBoolean started;
    private static final int MAX_SECOND = 60;
    public QpsLimit() {
        this.started = new AtomicBoolean(false);
        this.limit = Maps.newConcurrentMap();
        this.timer = new ScheduledThreadPoolExecutor(1, new ThreadFactory() {
            @Override
            public Thread newThread(Runnable r) {
                Thread t = new Thread(r);
                t.setName("tick-QpsLimit");
                return t;
            }
        });
    }

    /**
     * start qps limiter
     */
    public void start() {
        if (!started.compareAndSet(false, true)) {
            log.warn("QpsLimit already started!");
            return;
        }
        this.timer.scheduleAtFixedRate(new Runnable() {
            @Override
            public void run() {
                try {
                        limit.forEach((k, ref) -> {
                            if (ref != null) {
                                ref.get().get(Calendar.getInstance().get(Calendar.SECOND)).set(0L);
                            }
                        });
                } catch (Exception e) {
                    log.error("", e);
                }
            }
        }, 0, 1L, TimeUnit.SECONDS);
        log.info("QpsLimit start success!");
    }

    /**
     *
     * @param key
     * @param limit
     * @return if ret > 0 pass, else limit
     */
    public long limit(String key, int limit) {
        Assert.isTrue(this.started.get(), "QpsLimit should be invoke start method at first!");
        if (limit < 0) {
            return -1;
        }
        if (!this.limit.containsKey(key)) {
            if (this.limit.putIfAbsent(key, new SoftReference<>(Maps.newConcurrentMap())) == null) {
                synchronized (this.limit.get(key)) {
                    this.initQpsMap(this.limit.get(key).get());
                }
            }
        }
        return limit - this.limit.get(key).get().get(Calendar.getInstance().get(Calendar.SECOND)).incrementAndGet();
    }

    private void initQpsMap(Map<Integer, AtomicLong> m) {
        for (int i = 0; i < MAX_SECOND; i++) {
            m.put(i, new AtomicLong(0));
        }
    }

    /**
     * stop qps limiter
     */
    public void stop() {
        try {
            if (this.timer != null) {
                this.timer.shutdownNow();
            }
            this.limit.clear();
        } catch (Exception e) {
            log.error("", e);
        }
        log.info("QpsLimit stop!");
    }
}